# ------------------------------------------------------------------------
# Copyright (c) 2021 megvii-model. All Rights Reserved.
# ------------------------------------------------------------------------
# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# ------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------
"""
AnchorDETR model and criterion classes.
"""
import torch
from torch import nn

from util.misc import (NestedTensor, nested_tensor_from_tensor_list)

from .backbones.backbone import build_backbone
from .necks.anchor_transformer import build_transformer
from .matcher.matcher import build_matcher
from .losses.setcriterion import SetCriterion
from .postprocess.postprocess import PostProcess


def build(args):
    device = torch.device(args.device)
    num_classes = 20 if args.dataset_file != 'coco' else 91
    if args.dataset_file == "coco_panoptic":
        num_classes = 250
    backbone = build_backbone(args)
    transformer = build_transformer(args)
    model = AnchorDETR(
        backbone,
        transformer,
        num_feature_levels=args.num_feature_levels,
        aux_loss=args.aux_loss
    )
    # if args.masks:
    #     model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
    matcher = build_matcher(args, 'focal_loss')
    weight_dict = get_weight_dict(args)
    losses = get_losses(args)
    
    # num_classes, matcher, weight_dict, losses, focal_alpha=0.25
    criterion = SetCriterion(num_classes, matcher, weight_dict, losses, 
                             focal_alpha=args.focal_alpha, trans_type='ddetr')
    criterion.to(device)
    postprocessors = {'bbox': PostProcess(trans_type='ddetr')}
    # if args.masks:
    #     postprocessors['segm'] = PostProcessSegm()
    #     if args.dataset_file == "coco_panoptic":
    #         is_thing_map = {i: i <= 90 for i in range(201)}
    #         postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)

    return model, criterion, postprocessors


class AnchorDETR(nn.Module):
    """ This is the AnchorDETR module that performs object detection """

    def __init__(self, backbone, transformer,  num_feature_levels, aux_loss=True):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the transformer architecture. See transformer.py
            num_classes: number of object classes
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super().__init__()
        self.transformer = transformer
        hidden_dim = transformer.d_model
        self.backbone = backbone
        self.aux_loss = aux_loss

        #   摘抄自ddetr
        self.num_feature_levels = num_feature_levels
        if num_feature_levels > 1:
            num_backbone_outs = len(backbone.strides)
            input_proj_list = []
            for _ in range(num_backbone_outs):
                in_channels = backbone.num_channels[_]
                if _ == 0:
                    input_proj_list.append(nn.Sequential(
                        nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1),
                        nn.GroupNorm(32, hidden_dim),
                    ))
                else:
                    input_proj_list.append(nn.Sequential(
                        nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
                        nn.GroupNorm(32, hidden_dim),
                    ))
            self.input_proj = nn.ModuleList(input_proj_list)
        else:
            self.input_proj = nn.ModuleList([
                nn.Sequential(
                    nn.Conv2d(backbone.num_channels[0], hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                )])
        
        for proj in self.input_proj:
            nn.init.xavier_uniform_(proj[0].weight, gain=1)
            nn.init.constant_(proj[0].bias, 0)




    def forward(self, samples: NestedTensor):
        srcs, masks = self.backbone_forward(samples)
        outputs = self.neck_forward(srcs, masks) # hs, memory
        out = self.head_forward(outputs)
        out.update(outputs)
        return out



    def backbone_forward(self, samples: NestedTensor):
        if not isinstance(samples, NestedTensor):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)

        srcs = []
        masks = []
        for j, feat in enumerate(features):
            #if j== len(features)-1:
                src, mask = feat.decompose()
                srcs.append(self.input_proj(src))
                masks.append(mask)
                assert mask is not None
        srcs = torch.cat(srcs, dim=1)
        return srcs, masks
    
    def neck_forward(self, srcs, masks):
        outputs = self.transformer(srcs, masks)
        
        return outputs

    def head_forward(self, outputs):
        outputs_class = outputs['outputs_class']
        outputs_coord = outputs['outputs_coord']
        out = {'pred_logits': outputs_class[-1], 
               'pred_boxes': outputs_coord[-1]}
        if self.aux_loss and self.training:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        return out

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 
                 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

def get_weight_dict(args):
    weight_dict = {'loss_ce': 1, 
                   'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef
    if args.masks:
        weight_dict["loss_mask"] = args.mask_loss_coef
        weight_dict["loss_dice"] = args.dice_loss_coef
    # TODO this is a hack
    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)
    return weight_dict

def get_losses(args):
    losses = ['labels', 'boxes', 'cardinality']
    if args.masks:
        losses += ["masks"]
    return losses